from server import ws_connect
import asyncio
import psutil
from fastapi import WebSocket


async def get_cpu_usage():  # 获取cpu使用率
    loop = asyncio.get_event_loop()
    cpu_usage = await loop.run_in_executor(None, psutil.cpu_percent)
    return cpu_usage


async def get_memory_usage():  # 获取内存使用率
    loop = asyncio.get_event_loop()
    memory_info = await loop.run_in_executor(None, psutil.virtual_memory)
    memory_usage = memory_info.percent
    return memory_usage


async def get_disk_usage():  # 获取磁盘使用率
    loop = asyncio.get_event_loop()
    partitions = await loop.run_in_executor(None, psutil.disk_partitions)
    disk_usage_list = []

    for partition in partitions:  # 遍历所有分区
        try:
            usage = await loop.run_in_executor(None, psutil.disk_usage, partition.mountpoint)
            disk_usage_info = {
                'mountpoint': partition.mountpoint,
                'fstype': partition.fstype,
                'percent': usage.percent
            }
            disk_usage_list.append(disk_usage_info)
        except PermissionError:
            disk_usage_info = {
                'mountpoint': partition.mountpoint,
                'fstype': "unknown",
                'percent': 0
            }

    return disk_usage_list


async def get_system_info():  # 获取系统信息
    cpu_usage = await get_cpu_usage()
    memory_usage = await get_memory_usage()
    disk_usage = await get_disk_usage()

    system_info = {
        'cpu_usage': cpu_usage,
        'memory_usage': memory_usage,
        'disk_usage': disk_usage
    }

    return system_info


async def send_task(ws: WebSocket):  # 发送消息
    while True:
        system_info = await get_system_info()
        await ws.send_json({
            "type": "monitor",
            "data": system_info
        })
        await asyncio.sleep(5)

ws_connect.ws_tasks.append(send_task)
